import os
import random
import subprocess
from ultralytics import YOLO
import pandas as pd
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
from pytorch_grad_cam import GradCAM, EigenCAM
from pytorch_grad_cam.utils.image import preprocess_image
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget

# Installing libraries
try:
    import pytorch_grad_cam
except ImportError:
    subprocess.run(["pip", "install", "pytorch-grad-cam"])

# Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using the device: {device}")

def get_latest_train_folder():
    base_path = "C:/Users/caruso/Desktop/YOLO_Project/agrilus"
    train_folders = sorted([f for f in os.listdir(base_path) if f.startswith("train")], key=lambda x: int(x.replace("train", "")) if x.replace("train", "").isdigit() else -1, reverse=True)
    return os.path.join(base_path, train_folders[0], "weights") if train_folders else None

def train_model():
    model = YOLO("yolo11x-cls.pt").to(device)
    model.train(
        data='C:/Users/caruso/Desktop/YOLO_Project/images/',
        epochs=100,
        batch=16,
        imgsz=640,
        patience=10,
        project="agrilus",
        flipud=0.5,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.2,
        hsv_h=0.015, hsv_s=0.7, hsv_v=0.4,
        weight_decay=0.0005  # Added to reduce overfitting
    )

def predict_model():
    weights_folder = get_latest_train_folder()
    if not weights_folder:
        print("❌ Error: No training folders found!")
        return
    best_model_path = os.path.join(weights_folder, "best.pt")
    if not os.path.exists(best_model_path):
        print(f"❌ Error: File {best_model_path} does not exist! Verify the path.")
        return
    model = YOLO(best_model_path).to(device)
    source = "C:/Users/caruso/Desktop/YOLO_Project/images/test_random"
    model.predict(source=source, save=True, imgsz=640, conf=0.5)

def generate_CAM():
    weights_folder = get_latest_train_folder()
    if not weights_folder:
        print("❌ Error: No training folders found!")
        return
    best_model_path = os.path.join(weights_folder, "best.pt")
    if not os.path.exists(best_model_path):
        print(f"❌ Error: File {best_model_path} does not exist! Verify the path.")
        return
    model = YOLO(best_model_path).to(device)
    
    # Select all images in the folder
    test_dir = "C:/Users/caruso/Desktop/YOLO_Project/images/test_random"
    img_files = [f for f in os.listdir(test_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    if not img_files:
        print("No images found in test folder_random!")
        return
    
    # Creazione dell'oggetto CAM
    target_layers = [model.model.model[-2]]
    cam = EigenCAM(model.model, target_layers)
    
    output_dir = "C:/Users/caruso/Desktop/YOLO_Project/CAM_results"
    os.makedirs(output_dir, exist_ok=True)
    
    for img_file in img_files:
        img_path = os.path.join(test_dir, img_file)
        
        # Image uploading
        img = cv2.imread(img_path, 1)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_tensor = preprocess_image(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        img_tensor = img_tensor.to(device)
        
        # Generating CAM
        grayscale_cam = cam(input_tensor=img_tensor, targets=[ClassifierOutputTarget(0)])[0]
        grayscale_cam = cv2.resize(grayscale_cam, (img.shape[1], img.shape[0]))
        
        # Overlay with the original image
        heatmap = cv2.applyColorMap(np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET)
        cam_result = cv2.addWeighted(img, 0.5, heatmap, 0.5, 0)
        
        # Saving CAM image with unique name
        output_path = os.path.join(output_dir, f"cam_{os.path.splitext(img_file)[0]}.jpg")
        cv2.imwrite(output_path, cv2.cvtColor(cam_result, cv2.COLOR_RGB2BGR))
        print(f"CAM saved in {output_path}")

def main():
    train_model()
    predict_model()
    generate_CAM()

if __name__ == '__main__':
    main()
